package middleware

import (
	"context"

	"gitee.com/mkwchecking/sr-core/storage/database"
	"gitee.com/mkwchecking/sr-core/yerr"
	"github.com/gin-gonic/gin"
)

const (
	DataSetMiddleKey = "dataSet"
)

func WithDataset(ctx context.Context) (gin.HandlerFunc, error) {
	dataset, ok := database.FromCtx(ctx)
	if !ok {
		return nil, yerr.ErrBizPluginNotFound.Error("dataset")
	}

	return func(ctx *gin.Context) {
		ctx.Set(DataSetMiddleKey, dataset)
		ctx.Next()
	}, nil
}

func GetDataSet(ctx *gin.Context) (database.IDataSet, error) {

	if dataset, ok := ctx.Get(DataSetMiddleKey); ok {
		if ds, ok := dataset.(database.IDataSet); ok {
			return ds, nil
		}
	}

	return nil, yerr.ErrBizPluginNotFound.Error("dataset")
}
